# -*- coding: utf-8 -*-
from pymongo import UpdateOne, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes, get_trading_dates
from util.erio_util import dual_ammount
from util.database import DB_CONN
import pandas as pd
import numpy as np
from util.erio_util import dualdate, dualdateb, dualcode, dualbasicdate, dualbasicamount
"""
实现市盈率因子的计算和保存
"""
'''
ROE:分子是净利润,分母是净资产
'''


class SP_TTM_Factor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='pe')

    def dualdate(self, a):

        # print (a)
        if len(a.strip()) < 10:
            month = '0' + a.split('-')[1]
            pata = a.split('-')
            pata[1] = month
            final = '-'.join(pata)
            return (final)
        else:
            return a

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        dm = DataModule()
        dates = get_trading_dates(begin_date, end_date)
        codes = ['601808']
        update_requests = []
        for code in codes:
            mainall = dm.get_stock_mainreport(code)
            for date in dates:
                for asa in mainall['report_date'].apply(dualdate):
                    if date < asa:
                        continue
                    else:
                        reportdat = asa
                    break
                AT = mainall[mainall['report_date'] == reportdat]['zzczzy'].astype(
                    'float').tolist()[0]
                update_requests.append(
                    UpdateOne(
                        {'code': code, 'date': date},
                        {'$set': {'code': code, 'date': date, 'AT': AT}}, upsert=True))

                print (update_requests)
                break
            break

SP_TTM_Factor().compute('2018-01-02', '2018-09-05')
